package graph;

import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.SortedSet;
import java.util.TreeSet;

import kb.KnowledgeBase;
import correlation.HyperAlertCorrelationNode;
import dom.HyperAlert;

/**
 * Builds a correlation graph based on the Hyper Alert Correlation Nodes
 * @author Alex
 */
public class CorrelationGraphGenerator {

	private boolean debug = false;
	private final String OUTPUT_FILE = "output/AttackScenario.dot";
	private List<HyperAlertCorrelationNode> nodes;
	private String graphToExport = "";
	private Map<String,String> verticesHash = new HashMap<String,String>();
	
	public CorrelationGraphGenerator(List<HyperAlertCorrelationNode> nodes){
		
		this.nodes = nodes;
	
		if(this.nodes.isEmpty()){
			this.nodes = buildTestData();
		}
	}
	
	/**
	 * Conductor
	 */
	public void run(){
		print("entered run() in CorrelationGraphGenerator");
		buildHyperAlertCorrelationGraph();
		prepareHyperAlertCorrelationGraph();
		printHyperAlertCorrelationGraph();
	}
	
	/**
	 * Creates the graph in the form of a string
	 * @param nodes
	 */
	private void buildHyperAlertCorrelationGraph(){
		print("entered buildHyperAlertCorrelationGraph() in CorrelationGraphGenerator");
		print("received " + nodes.size() + " HyperAlertCorrelationNodes with which to build the graph");
		// begins at a node with no prerequisites which in theory should be a root node
		for(HyperAlertCorrelationNode hyperAlertCorrNode : nodes){
			for(HyperAlertCorrelationNode consequenceNode : hyperAlertCorrNode.getConsequenceNodes()){
				addVertice(hyperAlertCorrNode,consequenceNode);
			}
		}
	}
	
	/**
	 * Adds a vertice to the graph if it does not already exist in the graph
	 * @param rootNodeName
	 * @param consequenceNodeName
	 */
	private void addVertice(HyperAlertCorrelationNode rootNode, HyperAlertCorrelationNode consequenceNode){
		String key = format(rootNode.getHyperAlertTypeName()) +
				format(rootNode.getAlert().getFact("sourceIP")) +
				format(consequenceNode.getHyperAlertTypeName() + "_" +
						consequenceNode.getAlert().getFact("sourceIP"));
		String value = format(rootNode.getHyperAlertTypeName()) + "_" +
				format(rootNode.getAlert().getFact("sourceIP")) + " -> " + 
						format(consequenceNode.getHyperAlertTypeName()) + "_" +
						format(consequenceNode.getAlert().getFact("sourceIP")) + ";";
		if(!verticesHash.containsKey(key)){
			verticesHash.put(key,value);
			print("Adding vertice '"+value+"'");
		}
		else{
			print("Duplicate vertice [key='" +key+"',value='"+ value + "']");
		}
	}
	
	/**
	 * formats the value according to the requirements of GraphViz
	 * @param value
	 * @return
	 */
	private String format(String value){
		return value.trim().replace(" ","_");
	}
	
	private void prepareHyperAlertCorrelationGraph(){
		StringBuilder graphSB = new StringBuilder();		
		SortedSet<String> keys = new TreeSet<String>(verticesHash.keySet());
		
		if(keys.isEmpty()){
			print("the graph is empty :(");
		}
		
		graphSB.append("digraph {\n\ngraph [rankdir=LR];\n\n");
		
		for (String key : keys) { 
			graphSB.append(verticesHash.get(key) + "\n");
		}
		graphToExport = graphSB.append("\n}").toString();
	}
	
	/**
	 * Handles printing the graph
	 */
	private void printHyperAlertCorrelationGraph(){
		print("entered printHyperAlertCorrelationGraph() in CorrelationGraphGenerator");
		
		PrintWriter writer = null;

		try{
			writer = new PrintWriter(OUTPUT_FILE);
			writer.print(graphToExport);
		}
		catch(IOException ex){
			ex.printStackTrace();
			print("Exception while exporting graph. Cause:" + ex.getMessage());
		}
		finally{
			writer.close();
		}
		
		print("graph exported to: " + OUTPUT_FILE);
	}
	
	/**
	 * For now.
	 * @return List<HyperAlertCorrelationNode> test nodes
	 */
	public List<HyperAlertCorrelationNode> buildTestData(){
		
		print("===========================================================================");
		print("WARNING: USING TEST DATA --- LIST OF HYPER ALERT CORRELATION NODES IS EMPTY");
		print("===========================================================================");
		
		KnowledgeBase kb = KnowledgeBase.getInstance();
		List<HyperAlertCorrelationNode> testNodes = new ArrayList<HyperAlertCorrelationNode>();
		
		HyperAlert testAlert1 = new HyperAlert(kb.getHyperAlertType(384),new Date(),111222333444L, 1111, 123123123123L, 1234);
		HyperAlert testAlert2 = new HyperAlert(kb.getHyperAlertType(408),new Date(),111222333444L, 1111, 123123123123L, 1234);
		HyperAlert testAlert3 = new HyperAlert(kb.getHyperAlertType(402),new Date(),111222333444L, 1111, 123123123123L, 1234);
		HyperAlert testAlert4 = new HyperAlert(kb.getHyperAlertType(2256),new Date(),111222333444L, 1111, 123123123123L, 1234);
		
		HyperAlert testAlert5 = new HyperAlert(kb.getHyperAlertType(585),new Date(),111222333444L, 1111, 123123123123L, 1234);
		HyperAlert testAlert6 = new HyperAlert(kb.getHyperAlertType(716),new Date(),111222333444L, 1111, 123123123123L, 1234);
		HyperAlert testAlert7 = new HyperAlert(kb.getHyperAlertType(718),new Date(),111222333444L, 1111, 123123123123L, 1234);
		HyperAlert testAlert8 = new HyperAlert(kb.getHyperAlertType(2256),new Date(),111222333444L, 1111, 123123123123L, 1234);
		
		HyperAlertCorrelationNode testHyperAlertCorrelationNode1 = new HyperAlertCorrelationNode(testAlert1);
		HyperAlertCorrelationNode testHyperAlertCorrelationNode2 = new HyperAlertCorrelationNode(testAlert2);
		HyperAlertCorrelationNode testHyperAlertCorrelationNode3 = new HyperAlertCorrelationNode(testAlert3);
		HyperAlertCorrelationNode testHyperAlertCorrelationNode4 = new HyperAlertCorrelationNode(testAlert4);
		
		HyperAlertCorrelationNode testHyperAlertCorrelationNode5 = new HyperAlertCorrelationNode(testAlert5);
		HyperAlertCorrelationNode testHyperAlertCorrelationNode6 = new HyperAlertCorrelationNode(testAlert6);
		HyperAlertCorrelationNode testHyperAlertCorrelationNode7 = new HyperAlertCorrelationNode(testAlert7);
		HyperAlertCorrelationNode testHyperAlertCorrelationNode8 = new HyperAlertCorrelationNode(testAlert8);
		
		// Adding the consequences
		testHyperAlertCorrelationNode3.addConsequenceNode(testHyperAlertCorrelationNode4);
		testHyperAlertCorrelationNode2.addConsequenceNode(testHyperAlertCorrelationNode3);
		testHyperAlertCorrelationNode1.addConsequenceNode(testHyperAlertCorrelationNode2);
		//////////////////////////////
		testHyperAlertCorrelationNode7.addConsequenceNode(testHyperAlertCorrelationNode8);
		testHyperAlertCorrelationNode6.addConsequenceNode(testHyperAlertCorrelationNode7);
		testHyperAlertCorrelationNode5.addConsequenceNode(testHyperAlertCorrelationNode6);
		
		// Add nodes to list and return them
		testNodes.add(testHyperAlertCorrelationNode1);
		testNodes.add(testHyperAlertCorrelationNode2);
		testNodes.add(testHyperAlertCorrelationNode3);
		testNodes.add(testHyperAlertCorrelationNode4);
		testNodes.add(testHyperAlertCorrelationNode5);
		testNodes.add(testHyperAlertCorrelationNode6);
		testNodes.add(testHyperAlertCorrelationNode7);
		testNodes.add(testHyperAlertCorrelationNode8);
		
		return testNodes;
	}
	
	public void print(String value){
		if(debug){
			System.out.println(value);
		}
	}
	
}
